/* genlib.hpp
 * A testcase generator for Luogu OJ
 * License: MIT
 * sam@2aliens.net
 * */
#include "lib/indicators.hpp"
#include "lib/json.hpp"
#include "lib/templates.hpp"
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <chrono>
#include <thread>
#include <random>
#include <algorithm>
#include <unordered_set>
#include <unordered_map>
#include <limits.h>

#ifdef WIN32
#include <direct.h>
#include <io.h>
#else
#include <unistd.h>
#include <dirent.h>
#include <sys/stat.h>
#include <sys/types.h>
#endif

// Base class for testcase-files
class testcase {
    public:
        unsigned LEN;
        std::string dir;

        std::vector<nlohmann::json> data;

        testcase(const std::string &NAME, const unsigned &n) : LEN(n), dir("data") {
            data.resize(n + 1);

            if (access(dir.c_str(), 0) == -1) {
#ifdef WIN32
                int flag = mkdir(dir.c_str());
#else
                int flag = mkdir(dir.c_str(), S_IRWXU);
#endif
                if (flag == 0) {
                    std::cout << "Create directory successfully!" << "\n";

                } else {
                    std::cout << "Fail to create directory!" << "\n";
                    exit(0);
                }
            } else {
                std::cout << "This directory already exists." << "\n";
            }

            for (unsigned id = 1; id <= n; id++) {
                std::string serial = id < 10 ? "0" + std::to_string(id) : std::to_string(id);

                data[id]["inf"] = NAME + serial + ".in";
                data[id]["ouf"] = NAME + serial + ".out";
                data[id]["config"] = {
                    {"timeLimit", 1000},
                    {"isPretest", false},
                    {"memoryLimit", 128000},
                    {"score", 5}
                };
            }
        }

        void generate(void (*func) (unsigned));
        void solve(void (*func) ());
        void config();

        void set_timelimit(const unsigned &v, unsigned s = 1, unsigned t = 0);
        void set_memorylimit(const unsigned &v, unsigned s = 1, unsigned t = 0);
        void set_score(const unsigned &v, unsigned s = 1, unsigned t = 0);
        void set_ispretest(const bool &v, unsigned s = 1, unsigned t = 0);
        void set_subtaskid(const unsigned &v, unsigned s = 1, unsigned t = 0);
};

// Generate input files
void testcase::generate(void (*func) (unsigned)) {
    using namespace indicators;

    show_console_cursor(false);

    ProgressBar bar {
        option::BarWidth{50},
            option::Start{"["},
            option::Fill{"■"},
            option::Lead{"■"},
            option::Remainder{"-"},
            option::End{" ]"},
            option::MaxProgress{LEN},
            option::ForegroundColor{Color::cyan},
            option::FontStyles{std::vector<FontStyle>{FontStyle::bold}}
    };

    using namespace std::chrono;

    auto start = system_clock::now();
    for (unsigned id = 1; id <= LEN; id++) {
        bar.set_option(option::PostfixText {
                "Generating " + std::to_string(id) + "/" + std::to_string(LEN)
                });
        // this_thread::sleep_for(milliseconds(50));
        std::ofstream ofs(dir + "/" + data[id]["inf"].get<std::string>());
        std::streambuf *fout = std::cout.rdbuf(ofs.rdbuf());

        func(id);

        std::cout.rdbuf(fout);
        ofs.close();
        if (id == LEN) {
            auto end = system_clock::now();
            auto dur = duration_cast<milliseconds> (end - start);
            bar.set_option(option::PostfixText {
                    "Finished! " + std::to_string(dur.count()) + " ms"
                    });
        }
        bar.tick();
    }

    show_console_cursor(true);
}

// Generate config.yml
void testcase::config() {
    std::ofstream ofs(dir + "/config.yml");
    std::streambuf *fout = std::cout.rdbuf(ofs.rdbuf());

    for (unsigned id = 1; id <= LEN; id++) {
        nlohmann::json j = data[id];
        std::cout << j["inf"].get<std::string>() << ':' << "\n";
        for (auto &el : j["config"].items()) {
            std::cout << ' ' << el.key() << ": " << el.value() << "\n";
        }
        std::cout << "\n";
    }

    std::cout.rdbuf(fout);
    ofs.close();
}

// Generate answer files
void testcase::solve(void (*func) ()) {
    using namespace indicators;

    show_console_cursor(false);

    ProgressBar bar{
        option::BarWidth{50},
            option::Start{"["},
            option::Fill{"■"},
            option::Lead{"■"},
            option::Remainder{"-"},
            option::End{" ]"},
            option::MaxProgress{LEN},
            option::ForegroundColor{Color::cyan},
            option::FontStyles{std::vector<FontStyle>{FontStyle::bold}}
    };

    using namespace std::chrono;

    auto start = system_clock::now();
    for (unsigned id = 1; id <= LEN; id++) {
        bar.set_option(option::PostfixText {
                "Solving " + std::to_string(id) + "/" + std::to_string(LEN)
                });
        // this_thread::sleep_for(milliseconds(50));
        std::ifstream ifs(dir + "/" + data[id]["inf"].get<std::string>());
        std::ofstream ofs(dir + "/" + data[id]["ouf"].get<std::string>());
        std::streambuf *fin = std::cin.rdbuf(ifs.rdbuf());
        std::streambuf *fout = std::cout.rdbuf(ofs.rdbuf());

        func();

        std::cin.rdbuf(fin);
        std::cout.rdbuf(fout);
        ifs.close();
        ofs.close();
        if (id == LEN) {
            auto end = system_clock::now();
            auto dur = duration_cast<milliseconds> (end - start);
            bar.set_option(option::PostfixText {
                    "Finished! " + std::to_string(dur.count()) + " ms"
                    });
        }
        bar.tick();
    }

    show_console_cursor(true);
}

// Set timeLimit of testcases from No. s to t
void testcase::set_timelimit(const unsigned &v, unsigned s, unsigned t) {
    try {
        if (s > LEN || t > LEN)
            throw "Error (testcase::set_timelimit): Invalid range of testcases.";

        if (!t) t = LEN;

        for (unsigned id = s; id <= t; id++) data[id]["config"]["timeLimit"] = v;
    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}

// Set memoryLimit of testcases from No. s to t
void testcase::set_memorylimit(const unsigned &v, unsigned s, unsigned t) {
    try {
        if (s > LEN || t > LEN)
            throw "Error (testcase::set_memorylimit): Invalid range of testcases.";
        if (v > 516000)
            throw "Error (testcase::set_memorylimit): Memory limit should less than 516 MB";

        if (!t) t = LEN;

        for (unsigned id = s; id <= t; id++) data[id]["config"]["memoryLimit"] = v;
    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}

// Set score of testcases from No. s to t
void testcase::set_score(const unsigned &v, unsigned s, unsigned t) {
    try {
        if (s > LEN || t > LEN)
            throw "Error (testcase::set_score): Invalid range of testcases.";

        if (!t) t = LEN;

        for (unsigned id = s; id <= t; id++) data[id]["config"]["score"] = v;
    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}

// Set subtaskId of testcases from No. s to t
void testcase::set_subtaskid(const unsigned &v, unsigned s, unsigned t) {
    try {
        if (s > LEN || t > LEN)
            throw "Error (testcase::set_subtaskid): Invalid range of testcases.";

        if (!t) t = LEN;

        for (unsigned id = s; id <= t; id++) data[id]["config"]["subtaskId"] = v;
    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}

// Set isPretest of testcases from No. s to t
void testcase::set_ispretest(const bool &v, unsigned s, unsigned t) {
    try {
        if (s > LEN || t > LEN)
            throw "Error (testcase::set_ispretest): Invalid range of testcases.";

        if (!t) t = LEN;

        for (unsigned id = s; id <= t; id++) data[id]["config"]["isPretest"] = v;
    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}

// Base class of random seed
class rnd {
    protected:
        std::random_device rd;
        std::mt19937 _engine;

    public:
        rnd() : _engine(std::mt19937 {rd()}) {}

        static std::vector<int> permute(const size_t &n);
};

// Class of random integers
template <class T>
class integer : public rnd {
    private:
        std::uniform_int_distribution<T> _default;

    public:
        integer(const T &l = INT_MIN, const T &r = INT_MAX) : 
            _default(std::uniform_int_distribution<T> (l, r)) {}

        T tick() { return _default(_engine); };
};

// class of random floating-point values
template <class T>
class real : public rnd {
    private:
        std::uniform_real_distribution<T> _default;

    public:
        real(const T &l, const T &r) : _default(std::uniform_real_distribution (l, r)) {}

        T tick() { return _default(_engine); };
};

// Class of random char from a customized char-set
class charset : public rnd {
    private:
        std::uniform_int_distribution<> _default;
        std::vector<char> _charset;
        std::unordered_set<char> _set;

    public:
        charset(const std::string &s = "abcdefghijklmnopqrstuvwxyz") {
            for (const auto &c : s) {
                _set.insert(c);
            }
            _charset = std::vector<char>(_set.begin(), _set.end());
            _default = std::uniform_int_distribution<> (0, _charset.size() - 1);
        }

        char tick() { return _charset[_default(_engine)]; };

        std::string str(const size_t &len) {
            std::string res = "";
            for (size_t i = 1; i <= len; i++) {
                res += tick();
            }
            return res;
        }
};

// get random permutation from 1 to n
std::vector<int> rnd::permute(const size_t &n) {
    std::random_device rd;
    std::vector<int> res(n);
    for (size_t i = 0; i < n; i++) {
        res[i] = i + 1;
    }
    auto rnd_pos = [&](size_t s, size_t t) {return rd() % (t - s) + s;};
    for (size_t i = 0; i < n; i++) {
        std::swap(res[i], res[rnd_pos(i, n)]);
    }
    return res;
}

// Hash function of pair<int, int>
namespace std {
    template <> struct hash<std::pair<int, int>> {
        inline size_t operator()(const std::pair<int, int> &v) const {
            std::hash<int> int_hasher;
            return int_hasher(v.first) ^ int_hasher(v.second);
        }
    };
}

#define _MAX_EDGE (n * (n - 1)) >> 1
#define LL long long 

namespace graph {
    void tree(const unsigned &n, int l = 0, int r = 0);
    void connected_graph(const LL &n, const LL &m, int l = 0, int r = 0);
    void dag(const LL &n, const LL &m, int l = 0, int r = 0);
};

// get random tree of n vertices
// 4979 ms of 1e6 vertices
void graph::tree(const unsigned &n, int l, int r) {
    try {
        if (n > 1e6) 
            throw "Error (graph::tree): Number of vertices should less than 1e6.";

        unsigned m = 0;
        disjoint_set uf(n);

        integer<int> _vertex(1, n);
        integer<int> _weight(l, r);

        while (m < n - 1) {
            int u = 0, v = 0;
            while (u == v) {
                u = _vertex.tick();
                v = _vertex.tick();
            }
            if (uf.merge(u, v)) {
                m++;
                std::cout << u << ' ' << v;
                if (l || r) std::cout << ' ' << _weight.tick();
                std::cout << "\n";
            }
        }
    } catch (const char* msg){
        std::cerr << msg << "\n";
        exit(0);
    }
}

// get random connected graph of n vertices and m edges
// 5496 ms of 1e5 vertices and 1e6 edges
void graph::connected_graph(const LL &n, const LL &m, int l, int r) {
    try {
        if (n > 2e5) 
            throw "Error (graph::connected_graph): Number of vertices should less than 2e5.";
        if (m > 1e6) 
            throw "Error (graph::connected_graph): Number of edges should less than 1e6.";
        if (m < n || m > _MAX_EDGE)
            throw "Error (graph::connected_graph): Invalid number of vertices or edges.";

        disjoint_set uf(n);

        integer<int> _vertex(1, n);
        integer<int> _weight(l, r);

        std::unordered_set<std::pair<int, int>> st;

        unsigned cnt = 0;
        while (cnt < n - 1) {
            int u = 0, v = 0;
            while (u == v) {
                u = _vertex.tick();
                v = _vertex.tick();
            }
            if (uf.merge(u, v)) {
                cnt++;
                std::cout << u << ' ' << v;
                if (l || r) std::cout << ' ' << _weight.tick();
                std::cout << "\n";
                st.emplace(std::pair(std::min(u, v), std::max(u, v)));
            }
        }

        while (cnt < m) {
            int u = 0, v = 0;
            while (u == v) {
                u = _vertex.tick();
                v = _vertex.tick();
            }
            if (st.count(std::pair(std::min(u, v), std::max(u, v)))) continue;

            cnt++;
            std::cout << u << ' ' << v;
            if (l || r) std::cout << ' ' << _weight.tick();
            std::cout << "\n";
            st.emplace(std::pair(std::min(u, v), std::max(u, v)));
        }

    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}

// get random DAG of n vertices and m edges
// 6340 ms of 1e5 vertices and 1e6 edges
void graph::dag(const LL &n, const LL &m, int l, int r) {
    try {
        if (n > 2e5) 
            throw "Error (graph::dag): Number of vertices should less than 2e5.";
        if (m > 1e6) 
            throw "Error (graph::dag): Number of edges should less than 1e6.";
        if (m < n - 1 || m > _MAX_EDGE)
            throw "Error (graph::dag): Invalid number of vertices or edges.";

        disjoint_set uf(n);

        integer<int> _vertex(1, n);
        integer<int> _weight(l, r);

        std::unordered_set<std::pair<int, int>> st;
        std::unordered_map<int, int> mp;

        std::vector<int> topo_seq = rnd::permute(n);
        for (size_t i = 0; i < topo_seq.size(); i++) {
            mp.emplace(topo_seq[i], i);
        }

        unsigned cnt = 0;
        while (cnt < n - 1) {
            int u = 0, v = 0;
            while (u == v) {
                u = _vertex.tick();
                v = _vertex.tick();
            }
            if (uf.merge(u, v)) {
                cnt++;
                if (mp[u] > mp[v]) std::swap(u, v);
                std::cout << u << ' ' << v;
                if (l || r) std::cout << ' ' << _weight.tick();
                std::cout << "\n";
                st.emplace(std::pair(std::min(u, v), std::max(u, v)));
            }
        }

        while (cnt < m) {
            int u = 0, v = 0;
            while (u == v) {
                u = _vertex.tick();
                v = _vertex.tick();
            }
            if (st.count(std::pair(std::min(u, v), std::max(u, v)))) continue;

            cnt++;
            if (mp[u] > mp[v]) std::swap(u, v);
            std::cout << u << ' ' << v;
            if (l || r) std::cout << ' ' << _weight.tick();
            std::cout << "\n";
            st.emplace(std::pair(std::min(u, v), std::max(u, v)));
        }

    } catch (const char* msg) {
        std::cerr << msg << "\n";
        exit(0);
    }
}
